"""适配器模块"""

from typing import Dict, Type, List, Optional

from .base import AdapterBackend

# 适配器注册表：将适配器名称映射到适配器类
_ADAPTER_REGISTRY: Dict[str, Type[AdapterBackend]] = {}


def register_adapter(name: str, adapter_class: Type[AdapterBackend], aliases: Optional[List[str]] = None):
    """
    注册适配器
    
    Args:
        name: 适配器名称
        adapter_class: 适配器类（必须继承 BaseAdapter）
        aliases: 可选的别名列表
    
    Raises:
        TypeError: 适配器类未继承 BaseAdapter
    
    示例:
        >>> from bench.adapters import register_adapter
        >>> register_adapter("my_adapter", MyAdapter, aliases=["my_alias"])
    """
    if not issubclass(adapter_class, AdapterBackend):
        raise TypeError(f"{adapter_class} must inherit from BaseAdapter")
    
    _ADAPTER_REGISTRY[name.lower()] = adapter_class
    
    # 注册别名
    if aliases:
        for alias in aliases:
            _ADAPTER_REGISTRY[alias.lower()] = adapter_class


def get_adapter(adapter_backend: str) -> AdapterBackend:
    """
    根据适配器名称创建适配器实例
    
    Args:
        adapter_backend: 适配器名称（如 "onnxrt", "tensorrt", "sophon_tpu"）
    
    Returns:
        BaseAdapter 实例
    
    Raises:
        ValueError: 适配器名称未注册
    
    示例:
        >>> adapter = get_adapter("onnxrt")
        >>> handle = adapter.prepare(model_spec)
    """
    adapter_backend_lower = adapter_backend.lower()
    
    if adapter_backend_lower not in _ADAPTER_REGISTRY:
        available = ", ".join(sorted(_ADAPTER_REGISTRY.keys()))
        raise ValueError(
            f"Unsupported adapter: '{adapter_backend}'\n"
            f"Available adapters: {available}\n"
            f"To add a new adapter, use register_adapter() or add it to "
            f"src/bench/adapters/__init__.py"
        )
    
    adapter_class = _ADAPTER_REGISTRY[adapter_backend_lower]
    return adapter_class()


def list_adapters(include_unavailable: bool = True) -> List[str]:
    """
    列出所有已注册的适配器
    
    Args:
        include_unavailable: 是否包含不可用的适配器
    
    Returns:
        适配器名称列表
    
    示例:
        >>> # 列出所有已注册的适配器
        >>> list_adapters(include_unavailable=True)
        ['onnxruntime', 'tensorrt', 'sophon_tpu']
        
        >>> # 只列出可用的适配器
        >>> list_adapters(include_unavailable=False)
        ['onnxruntime']
    """
    adapters = sorted(_ADAPTER_REGISTRY.keys())
    
    if include_unavailable:
        return adapters
    
    # 只返回可用的适配器
    available = []
    for name in adapters:
        adapter_class = _ADAPTER_REGISTRY[name]
        if adapter_class.is_available():
            available.append(name)
    
    return available


# 导入所有适配器以触发注册
from .onnxrt import OnnxRuntimeAdapter

# 注册所有适配器
register_adapter("onnxruntime", OnnxRuntimeAdapter, aliases=["onnxrt"])


# 公开的适配器类（供类型检查使用）
__all__ = [
    "AdapterBackend",
    "OnnxRuntimeAdapter",
    "InferenceRunner",
    "get_adapter",
    "list_adapters",
    "register_adapter",
]

